import os
from CalTorque import CalTorque
from datetime import datetime
from PlotCalTorque import PlotCalTorque 
from string import Template
import numpy
import matplotlib
matplotlib.use('agg')
from matplotlib import pyplot as PLT
import pickle
from collections import Counter
from decimal import Decimal

def SaveAveragePKL(x, y):
    avgdata = {}
    avgdata['x'] = x
    avgdata['y'] = y
    output = open('pickled/pickled_average_in.pkl', 'wb')
    pickle.dump(avgdata, output, -1)
    output.close()

def AverageOUT(data_files, bins, max_d):
    x_inputs = numpy.zeros(bins)
    X = numpy.linspace(0,max_d,bins)
    Y = numpy.zeros(bins)
    Y_std = numpy.zeros(bins)
    Y2 = []
    
    for i in range(X.size):
        Y2.append([])
    
    for f in data_files:
       if (f.split('.')[-1] not in ['xlsx', 'xls']):
          continue
       data = CalTorque(f)
       if data['direction'] != 'out':
          continue #Skip all insertions
       if data['source'] != 'weakTh-228':
          continue #Skip all other sources
       if data['station'] in ['S8', 'S11', 'Other']:
          continue
       print f
       
       for j,x in enumerate(X):
          pos_torque = []
          x0 = x-(max_d/bins)/2.
          x1 = x+(max_d/bins)/2.
          for i, p in enumerate(data['pos_offsetfix']):
             if p>x0 and p<=x1:
                pos_torque.append(data['torque_inoz'][i])
          if len(pos_torque)>0:
             if min(pos_torque) == 0: continue
             Y[j] += min(pos_torque)
             Y2[j].append(min(pos_torque))
             x_inputs[j]+=1.
    for i,n in enumerate(x_inputs):
       if n==0: continue
       Y[i] = Y[i]/n
       Y_std[i] = numpy.std(Y2[i], dtype = numpy.float64)
    
    #SaveAveragePKL(X, Y)
    
    fig = PLT.figure(figsize=(15,8),dpi=150)
    box = [0.14, 0.14, 0.76, 0.76]
    ax1 = fig.add_axes(box)
    ax1.set_xlabel('Position (cm)')
    ax1.set_ylabel('Torque (in.oz)')
    miny = min(Y_std)
    maxy = max(data['torque_inoz'])
    maxx = max(data['position'])
    
    ## Add station locations
    if maxx >= 328.9:
        ax1.axvline(338.9, color = 'r', linewidth = 1, linestyle = '-')
        ax1.annotate('S2', xy = (340.9, miny * 0.8), color = 'r')
    if maxx >= 382.6:
        ax1.axvline(392.6, color = 'r', linewidth = 1, linestyle = '-')
        ax1.annotate('S5', xy = (394.6, miny * 0.8), color = 'r')
    if maxx >= 435.7:
        ax1.axvline(445.7, color = 'r', linewidth = 1, linestyle = '-')
        ax1.annotate('S8', xy = (447.7, miny * 0.8), color = 'r')
    if maxx >= 491.7:
        ax1.axvline(501.7, color = 'r', linewidth = 1, linestyle = '-')
        ax1.annotate('S11', xy = (503.7, miny * 0.8), color = 'r')
    if maxx >= 595.8:
        ax1.axvline(605.8, color = 'r', linewidth = 1, linestyle = '-')
        ax1.annotate('S17', xy = (607.8, miny * 0.8), color = 'r')
    
    ## add bend locations
    if maxx >= 47.8:
        ax1.axvspan(47.8, 93.5, color = 'g', alpha = 0.2)
        ax1.annotate('Bellows', xy = (48, miny), color = 'g')
    if maxx >= 136.4:
        ax1.axvspan(136.4, 144.6, color = 'g', alpha = 0.2)
        ax1.annotate('First outer bend', xy = (136, miny), color = 'g')
    if maxx >= 196.5:
        ax1.axvspan(196.5, 218.0, color = 'g', alpha = 0.2)
        ax1.annotate('The "S"', xy = (212, miny), color = 'g')
    if maxx >= 228.3:
        ax1.axvspan(228.3, 241.5, color = 'g', alpha = 0.2)
    if maxx >= 274.6:
        ax1.axvspan(274.6, 284.6, color = 'g', alpha = 0.2)
        ax1.annotate('Bend A5', xy = (274, miny), color = 'g')
    if maxx >= 309.0:
        ax1.axvspan(309.0, 318.9, color = 'g', alpha = 0.2)
        ax1.annotate('Bend A6', xy = (309, miny), color = 'g')
    if maxx >= 358.8:
        ax1.axvspan(358.8, 370.8, color = 'g', alpha = 0.2)
        ax1.annotate('Bend A7', xy = (359, miny), color = 'g')
    if maxx >= 414.5:
        ax1.axvspan(414.5, 426.5, color = 'g', alpha = 0.2)
        ax1.annotate('Bend A8', xy = (414, miny), color = 'g')
    if maxx >= 445.3:
        ax1.axvspan(445.3, 455.3, color = 'g', alpha = 0.2)
        ax1.annotate('Bend A9', xy = (445, miny), color = 'g')
    if maxx >= 467.5:
        ax1.axvspan(467.5, 479.5, color = 'g', alpha = 0.2)
        ax1.annotate('Bend A10', xy = (467, miny), color = 'g')
    if maxx >= 523.8:
        ax1.axvspan(523.8, 535.8, color = 'g', alpha = 0.2)
        ax1.annotate('Bend A11', xy = (523, miny), color = 'g')
    if maxx >= 571.6:
        ax1.axvspan(571.6, 583.6, color = 'g', alpha = 0.2)
        ax1.annotate('Bend A12', xy = (571, miny), color = 'g')
    if maxx >= 627.9:
        ax1.axvspan(627.9, 639.9, color = 'g', alpha = 0.2)
        ax1.annotate('Bend A13', xy = (627, miny), color = 'g')
    if maxx >= 652.7:
        ax1.axvspan(652.7, 660.7, color = 'g', alpha = 0.2)
        ax1.annotate('Bend A14', xy = (652, miny), color = 'g')
    if maxx >= 678.5:
        ax1.axvspan(678.5, 681.1, color = 'g', alpha = 0.2)
        ax1.annotate('Bend A15', xy = (678, miny), color = 'g')
    
    PLT.title('Average of all OUT deployments')
    ax1.grid()
    # PLT.plot(X,Y)
    PLT.errorbar(X, Y, xerr = None, yerr = Y_std)
    fig.savefig('AveragePlots/average_out.png')

dp = 'data/'
ls = os.listdir(dp)
data_files = []
for l in ls:
    if l.split('.')[-1] not in ['xls', 'xlsx']: continue
    if 'out' not in l: continue
    data_files.append(dp + l)

AverageOUT(data_files, 400, 400)